// Copyright 2015 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef COMPONENTS_POLICY_CORE_COMMON_REMOTE_COMMANDS_TEST_SUPPORT_TESTING_REMOTE_COMMANDS_SERVER_H_
#define COMPONENTS_POLICY_CORE_COMMON_REMOTE_COMMANDS_TEST_SUPPORT_TESTING_REMOTE_COMMANDS_SERVER_H_

#include <stddef.h>

#include <memory>
#include <string>
#include <vector>

#include "base/functional/callback_forward.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/weak_ptr.h"
#include "base/synchronization/lock.h"
#include "base/threading/thread_checker.h"
#include "components/policy/core/common/remote_commands/remote_command_job.h"
#include "components/policy/proto/device_management_backend.pb.h"

namespace base {
class TickClock;
class SingleThreadTaskRunner;
}  // namespace base

namespace policy {

using SignatureType = enterprise_management::PolicyFetchRequest::SignatureType;

// Callback called when a command's result is reported back to the server.
using ResultReportedCallback =
    base::OnceCallback<void(const enterprise_management::RemoteCommandResult&)>;

std::string SignDataWithTestKey(const std::string& data,
                                SignatureType algorithm);

// This class implements server-side logic for remote commands service tests. It
// acts just like a queue, and there are mainly two exposed methods for this
// purpose. Test authors are expected to call IssueCommand() to add commands to
// the queue in their tests. FetchCommands() should be called when serving a
// request to the server intercepted e.g. via a net::URLRequestInterceptor or a
// mock CloudPolicyClient.
// In addition, this class also supports verifying command result sent back to
// server, and test authors should provide a callback for this purpose when
// issuing a command.
// All methods except FetchCommands() should be called from the same thread, and
// FetchCommands() can be called from any thread.
// Note that test author is responsible for ensuring that FetchCommands() is not
// called from another thread after |this| has been destroyed.
class TestingRemoteCommandsServer {
 public:
  TestingRemoteCommandsServer();
  TestingRemoteCommandsServer(const TestingRemoteCommandsServer&) = delete;
  TestingRemoteCommandsServer& operator=(const TestingRemoteCommandsServer&) =
      delete;
  virtual ~TestingRemoteCommandsServer();

  using RemoteCommandResults =
      std::vector<enterprise_management::RemoteCommandResult>;

  using RemoteCommands = std::vector<enterprise_management::SignedData>;

  // Add a command that will be returned on the next FetchCommands() call.
  // |clock_| will be used to get the command issue time.
  // |reported_callback| will be called from the same thread when a command
  // result is reported back to the server.
  void IssueCommand(const enterprise_management::SignedData& signed_data,
                    ResultReportedCallback reported_callback);

  // Calling this will tell the server to return no commands on the next
  // call to FetchCommands().
  // All issued commands will instead be returned the second time
  // FetchCommands() is called.
  void OnNextFetchCommandsCallReturnNothing();

  // Fetch commands, acknowledging all commands up to and including
  // |last_command_id|, and provide |previous_job_results| as results for
  // previous commands. |last_command_id| can be a nullptr which indicates that
  // no commands will be acknowledged.
  // See the protobuf definition for a definition of the protocol between server
  // and client for remote command fetching.
  // Unlike every other methods in the class, this method can be called from
  // any thread.
  RemoteCommands FetchCommands(
      std::unique_ptr<RemoteCommandJob::UniqueIDType> last_command_id,
      const RemoteCommandResults& previous_job_results);

  // Set alternative clock for obtaining the command issue time. The default
  // clock uses the system clock.
  void SetClock(const base::TickClock* clock);

  // Get the number of commands for which no results have been reported yet.
  // This number also includes commands which have not been fetched yet.
  size_t NumberOfCommandsPendingResult() const;

  // The latest command ID generated by server.
  RemoteCommandJob::UniqueIDType last_command_id() const {
    return last_generated_unique_id_;
  }

  RemoteCommandJob::UniqueIDType GetNextCommandId() {
    return ++last_generated_unique_id_;
  }

 private:
  struct RemoteCommandWithCallback;

  void HandleRemoteCommandResults(const RemoteCommandResults& results);

  void ReportJobResult(
      ResultReportedCallback reported_callback,
      const enterprise_management::RemoteCommandResult& job_result) const;

  // If true no commands will be returned to the caller of the next
  // FetchCommands() call.
  bool return_nothing_for_next_fetch_ = false;

  // The main command queue.
  std::vector<RemoteCommandWithCallback> commands_;

  // The latest command ID generated by server. This test server generates
  // strictly monotonically increasing unique IDs.
  RemoteCommandJob::UniqueIDType last_generated_unique_id_ = 0;

  // Clock used to generate command issue time when IssueCommand() is called.
  raw_ptr<const base::TickClock> clock_;

  // A lock protecting the command queues, as well as generated and acknowledged
  // IDs.
  base::Lock lock_;

  // The thread on which this test server is created, usually the UI thread.
  scoped_refptr<base::SingleThreadTaskRunner> task_runner_;

  // A weak pointer created early so that it can be accessed from other threads.
  base::WeakPtr<TestingRemoteCommandsServer> weak_ptr_to_this_;

  base::ThreadChecker thread_checker_;
  base::WeakPtrFactory<TestingRemoteCommandsServer> weak_factory_{this};
};

}  // namespace policy

#endif  // COMPONENTS_POLICY_CORE_COMMON_REMOTE_COMMANDS_TEST_SUPPORT_TESTING_REMOTE_COMMANDS_SERVER_H_
